package controllers

import (
	"github.com/kevinstudy2021/rabbitmq-operator/api/v1beta1"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
)

var (
	RabbitmqLabelKey         = "rabbitmq.aliee.io/single"
	RabbitmqLabelValue       = "rabbitmq-cluster"
	RabbitmqCommonLabelKey   = "app"
	RabbitmqCommonLabelValue = "rabbitmq-cluster"
	RabbitmqDataDirName      = "datadir"
	True                     = true
	False                    = false
)

func MutateStatefulSet(cluster *v1beta1.Svc, sts *appsv1.StatefulSet) {
	//var sc string
	var localpath string
	//sc = "managed-nfs-storage"
	localpath = "local-storage"
	sts.Labels = map[string]string{
		RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
	}
	sts.Spec = appsv1.StatefulSetSpec{
		Replicas:    cluster.Size,
		ServiceName: cluster.Name,
		Selector: &metav1.LabelSelector{MatchLabels: map[string]string{
			RabbitmqLabelKey: cluster.Name,
		}},
		VolumeClaimTemplates: []corev1.PersistentVolumeClaim{
			corev1.PersistentVolumeClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name: RabbitmqDataDirName,
				},
				Spec: corev1.PersistentVolumeClaimSpec{
					AccessModes: []corev1.PersistentVolumeAccessMode{
						corev1.ReadWriteOnce,
					},
					//StorageClassName: &sc,
					StorageClassName: &localpath,
					Resources: corev1.ResourceRequirements{
						Requests: corev1.ResourceList{
							corev1.ResourceStorage: resource.MustParse("1Gi"),
						},
					},
				},
			},
		},
		Template: corev1.PodTemplateSpec{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					RabbitmqLabelKey:       cluster.Name,
					RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
				},
			},
			Spec: corev1.PodSpec{
				Containers:     newContainers(cluster),
				InitContainers: NewInitContainers(cluster),
			},
		},
	}
}

func MutateClusterStatefulset(cluster *v1beta1.Svc, sts *appsv1.StatefulSet) {
	//var sc string
	var localpath string
	//sc = "managed-nfs-storage"
	localpath = "local-storage"
	volume := corev1.Volume{
		Name: "rabbitmq-config-volume",
		VolumeSource: corev1.VolumeSource{
			//EmptyDir: &corev1.EmptyDirVolumeSource{},
			ConfigMap: &corev1.ConfigMapVolumeSource{
				LocalObjectReference: corev1.LocalObjectReference{
					Name: cluster.Name,
				},
				Items: []corev1.KeyToPath{
					{
						Key:  "rabbitmq.conf",
						Path: "rabbitmq.conf",
					},
					{
						Key:  "enabled_plugins",
						Path: "enabled_plugins",
					},
				},
			},
		},
	}
	sts.Labels = map[string]string{
		RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
	}
	sts.Spec = appsv1.StatefulSetSpec{
		Replicas:    cluster.Size,
		ServiceName: cluster.Name,
		Selector: &metav1.LabelSelector{MatchLabels: map[string]string{
			RabbitmqLabelKey: cluster.Name,
		}},
		VolumeClaimTemplates: []corev1.PersistentVolumeClaim{
			corev1.PersistentVolumeClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name: RabbitmqDataDirName,
				},
				Spec: corev1.PersistentVolumeClaimSpec{
					AccessModes: []corev1.PersistentVolumeAccessMode{
						corev1.ReadWriteOnce,
					},
					//StorageClassName: &sc,
					StorageClassName: &localpath,
					Resources: corev1.ResourceRequirements{
						Requests: corev1.ResourceList{
							corev1.ResourceStorage: resource.MustParse("1Gi"),
						},
					},
				},
			},
		},

		Template: corev1.PodTemplateSpec{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					RabbitmqLabelKey:       cluster.Name,
					RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
				},
			},
			Spec: corev1.PodSpec{
				Containers:         newContainersCluster(cluster),
				InitContainers:     NewInitContainers(cluster),
				Volumes:            []corev1.Volume{volume},
				ServiceAccountName: cluster.Name,
			},
		},
	}
}

func newContainers(cluster *v1beta1.Svc) []corev1.Container {
	return []corev1.Container{
		corev1.Container{
			Name:  "rabbitmq",
			Image: cluster.Image,
			Ports: []corev1.ContainerPort{
				corev1.ContainerPort{
					Name:          "tcp",
					ContainerPort: 5672,
				},
				corev1.ContainerPort{
					Name:          "web",
					ContainerPort: 15672,
				},
			},
			Env: []corev1.EnvVar{
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_USER",
					Value: cluster.Username,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_PASS",
					Value: cluster.Password,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_MNESIA_BASE",
					Value: "/var/lib/rabbitmq/mnesia/rabbit@rabbit-server",
				},
			},
			ImagePullPolicy: corev1.PullPolicy("IfNotPresent"),
			VolumeMounts: []corev1.VolumeMount{
				corev1.VolumeMount{
					Name:      RabbitmqDataDirName,
					MountPath: "/var/lib/rabbitmq/mnesia/rabbit@rabbit-server",
				},
			},
		},
	}
}

func newContainersCluster(cluster *v1beta1.Svc) []corev1.Container {

	return []corev1.Container{
		corev1.Container{
			Name:  "rabbitmq-cluster",
			Image: cluster.Image,
			Args: []string{
				"-c", "cp -v /etc/rabbitmq/rabbitmq.conf ${RABBITMQ_CONFIG_FILE}; exec docker-entrypoint.sh rabbitmq-server",
			},
			Command: []string{
				"sh",
			},
			Ports: []corev1.ContainerPort{
				corev1.ContainerPort{
					Name:          "tcp",
					ContainerPort: 5672,
				},
				corev1.ContainerPort{
					Name:          "web",
					ContainerPort: 15672,
				},
			},
			Env: []corev1.EnvVar{
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_USER",
					Value: cluster.Username,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_PASS",
					Value: cluster.Password,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_MNESIA_BASE",
					Value: "/var/lib/rabbitmq/mnesia/rabbit@rabbit-server",
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_ERLANG_COOKIE",
					Value: "RABBITMQ_COOKIE",
				},
				corev1.EnvVar{
					Name:  "K8S_SERVICE_NAME",
					Value: cluster.Name,
				},
				corev1.EnvVar{
					Name: "POD_IP",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "status.podIP",
						},
					},
				},
				corev1.EnvVar{
					Name: "POD_NAME",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "metadata.name",
						},
					},
				},
				corev1.EnvVar{
					Name: "POD_NAMESPACE",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "metadata.namespace",
						},
					},
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_USE_LONGNAME",
					Value: "true",
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_NODENAME",
					Value: "rabbit@$(POD_NAME).rabbitmq.$(POD_NAMESPACE).svc.cluster.local",
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_CONFIG_FILE",
					Value: "/var/lib/rabbitmq/rabbitmq.conf",
				},
			},
			LivenessProbe: &corev1.Probe{
				ProbeHandler: corev1.ProbeHandler{
					Exec: &corev1.ExecAction{
						Command: []string{"rabbitmqctl", "status"},
					},
					//TCPSocket: &corev1.TCPSocketAction{
					//	Port: intstr.IntOrString{
					//		IntVal: 5672,
					//	},
					//},
				},

				InitialDelaySeconds: 30,
				TimeoutSeconds:      10,
			},
			ReadinessProbe: &corev1.Probe{
				ProbeHandler: corev1.ProbeHandler{
					Exec: &corev1.ExecAction{
						Command: []string{"rabbitmqctl", "status"},
					},
					//TCPSocket: &corev1.TCPSocketAction{
					//	Port: intstr.IntOrString{
					//		IntVal: 5672,
					//	},
					//},
				},
				InitialDelaySeconds: 10,
				TimeoutSeconds:      10,
			},
			ImagePullPolicy: corev1.PullPolicy("IfNotPresent"),
			VolumeMounts: []corev1.VolumeMount{
				//corev1.VolumeMount{
				//	Name:      RabbitmqDataDirName,
				//	MountPath: "/var/lib/rabbitmq/mnesia/rabbit@rabbit-server",
				//},
				corev1.VolumeMount{
					Name:      "rabbitmq-config-volume",
					MountPath: "/etc/rabbitmq",
				},
			},
		},
	}
}

func NewInitContainers(cluster *v1beta1.Svc) []corev1.Container {
	return []corev1.Container{
		corev1.Container{
			Command: []string{
				"bin/sh", "-ec", "echo 10000 > /proc/sys/net/core/somaxconn",
			},
			Image:           "busybox",
			ImagePullPolicy: corev1.PullIfNotPresent,
			Name:            "init",
			SecurityContext: &corev1.SecurityContext{
				Privileged: &True,
			},
		},
	}
}

func newClusterContainers(cluster *v1beta1.Svc) []corev1.Container {
	return []corev1.Container{
		corev1.Container{
			Name:  "rabbitmq",
			Image: cluster.Image,
			Ports: []corev1.ContainerPort{
				corev1.ContainerPort{
					Name:          "tcp",
					ContainerPort: 5672,
				},
				corev1.ContainerPort{
					Name:          "web",
					ContainerPort: 15672,
				},
			},
			Env: []corev1.EnvVar{
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_USER",
					Value: cluster.Username,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_DEFAULT_PASS",
					Value: cluster.Password,
				},
				corev1.EnvVar{
					Name:  "RABBITMQ_MNESIA_BASE",
					Value: "/var/lib/rabbitmq/mnesia/rabbit@rabbit-server",
				},
				corev1.EnvVar{
					Name: "POD_IP",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "status.podIP",
						},
					},
				},
				corev1.EnvVar{
					Name: "MY_NAMESPACE",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "metadata.namespace",
						},
					},
				},
			},
			VolumeMounts: []corev1.VolumeMount{
				corev1.VolumeMount{
					Name:      RabbitmqDataDirName,
					MountPath: "/var/run/etcd",
				},
			},
			Command: []string{
				"/bin/sh", "-ec",
				"HOSTNAME=$(hostname)\n\n              ETCDCTL_API=3\n\n              eps() {\n                  EPS=\"\"\n                  for i in $(seq 0 $((${INITIAL_CLUSTER_SIZE} - 1))); do\n                      EPS=\"${EPS}${EPS:+,}http://${SET_NAME}-${i}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2379\"\n                  done\n                  echo ${EPS}\n              }\n\n              member_hash() {\n                  etcdctl member list | grep -w \"$HOSTNAME\" | awk '{ print $1}' | awk -F \",\" '{ print $1}'\n              }\n\n              initial_peers() {\n                  PEERS=\"\"\n                  for i in $(seq 0 $((${INITIAL_CLUSTER_SIZE} - 1))); do\n                    PEERS=\"${PEERS}${PEERS:+,}${SET_NAME}-${i}=http://${SET_NAME}-${i}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2380\"\n                  done\n                  echo ${PEERS}\n              }\n\n              # etcd-SET_ID\n              SET_ID=${HOSTNAME##*-}\n\n              # adding a new member to existing cluster (assuming all initial pods are available)\n              if [ \"${SET_ID}\" -ge ${INITIAL_CLUSTER_SIZE} ]; then\n                  # export ETCDCTL_ENDPOINTS=$(eps)\n                  # member already added?\n\n                  MEMBER_HASH=$(member_hash)\n                  if [ -n \"${MEMBER_HASH}\" ]; then\n                      # the member hash exists but for some reason etcd failed\n                      # as the datadir has not be created, we can remove the member\n                      # and retrieve new hash\n                      echo \"Remove member ${MEMBER_HASH}\"\n                      etcdctl --endpoints=$(eps) member remove ${MEMBER_HASH}\n                  fi\n\n                  echo \"Adding new member\"\n\n                  etcdctl member --endpoints=$(eps) add ${HOSTNAME} --peer-urls=http://${HOSTNAME}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2380 | grep \"^ETCD_\" > /var/run/etcd/new_member_envs\n\n                  if [ $? -ne 0 ]; then\n                      echo \"member add ${HOSTNAME} error.\"\n                      rm -f /var/run/etcd/new_member_envs\n                      exit 1\n                  fi\n\n                  echo \"==> Loading env vars of existing cluster...\"\n                  sed -ie \"s/^/export /\" /var/run/etcd/new_member_envs\n                  cat /var/run/etcd/new_member_envs\n                  . /var/run/etcd/new_member_envs\n\n                  exec etcd --listen-peer-urls http://${POD_IP}:2380 \\\n                      --listen-client-urls http://${POD_IP}:2379,http://127.0.0.1:2379 \\\n                      --advertise-client-urls http://${HOSTNAME}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2379 \\\n                      --data-dir /var/run/etcd/default.etcd\n              fi\n\n              for i in $(seq 0 $((${INITIAL_CLUSTER_SIZE} - 1))); do\n                  while true; do\n                      echo \"Waiting for ${SET_NAME}-${i}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local to come up\"\n                      ping -W 1 -c 1 ${SET_NAME}-${i}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local > /dev/null && break\n                      sleep 1s\n                  done\n              done\n\n              echo \"join member ${HOSTNAME}\"\n              # join member\n              exec etcd --name ${HOSTNAME} \\\n                  --initial-advertise-peer-urls http://${HOSTNAME}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2380 \\\n                  --listen-peer-urls http://${POD_IP}:2380 \\\n                  --listen-client-urls http://${POD_IP}:2379,http://127.0.0.1:2379 \\\n                  --advertise-client-urls http://${HOSTNAME}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2379 \\\n                  --initial-cluster-token etcd-cluster-1 \\\n                  --data-dir /var/run/etcd/default.etcd \\\n                  --initial-cluster $(initial_peers) \\\n                  --initial-cluster-state new",
			},
			Lifecycle: &corev1.Lifecycle{
				PreStop: &corev1.LifecycleHandler{
					Exec: &corev1.ExecAction{
						Command: []string{
							"/bin/sh", "-ec",
							"HOSTNAME=$(hostname)\n\n                    member_hash() {\n                        etcdctl member list | grep -w \"$HOSTNAME\" | awk '{ print $1}' | awk -F \",\" '{ print $1}'\n                    }\n\n                    eps() {\n                        EPS=\"\"\n                        for i in $(seq 0 $((${INITIAL_CLUSTER_SIZE} - 1))); do\n                            EPS=\"${EPS}${EPS:+,}http://${SET_NAME}-${i}.${SET_NAME}.${MY_NAMESPACE}.svc.cluster.local:2379\"\n                        done\n                        echo ${EPS}\n                    }\n\n                    export ETCDCTL_ENDPOINTS=$(eps)\n                    SET_ID=${HOSTNAME##*-}\n\n                    # Removing member from cluster\n                    if [ \"${SET_ID}\" -ge ${INITIAL_CLUSTER_SIZE} ]; then\n                        echo \"Removing ${HOSTNAME} from etcd cluster\"\n                        etcdctl member remove $(member_hash)\n                        if [ $? -eq 0 ]; then\n                            # Remove everything otherwise the cluster will no longer scale-up\n                            rm -rf /var/run/etcd/*\n                        fi\n                    fi",
						},
					},
				},
			},
		},
	}
}

//func MutateServiceAccount(cluster *v1beta1.Svc, sa *corev1.ServiceAccount) {
//	sa.Name = cluster.Name
//	sa.Namespace = cluster.Namespace
//}

func MutateRole(cluster *v1beta1.Svc, role *v1.Role) {
	role.Rules = []v1.PolicyRule{
		{
			APIGroups: []string{""},
			Resources: []string{"endpoints"},
			Verbs:     []string{"get"},
		},
	}
}

func MutateRoleBinding(cluster *v1beta1.Svc, rb *v1.RoleBinding, rabbitmq *v1beta1.Rabbitmq) {
	rb.RoleRef = v1.RoleRef{
		APIGroup: "rbac.authorization.k8s.io",
		Kind:     "Role",
		Name:     cluster.Name,
	}
	rb.Subjects = []v1.Subject{
		{
			Kind:      "ServiceAccount",
			Name:      cluster.Name,
			Namespace: rabbitmq.Namespace,
		},
	}
}

func MutateHeadlessSvc(cluster *v1beta1.Svc, svc *corev1.Service) {
	svc.Labels = map[string]string{
		RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
	}
	svc.Spec = corev1.ServiceSpec{
		ClusterIP: corev1.ClusterIPNone,
		Selector: map[string]string{
			RabbitmqLabelKey: cluster.Name,
		},
		Ports: []corev1.ServicePort{
			corev1.ServicePort{
				Name: "tcp",
				Port: 5672,
			},
			corev1.ServicePort{
				Name: "web",
				Port: 15672,
			},
		},
	}
}

func MutateNodePortSvc(cluster *v1beta1.Svc, svc *corev1.Service) {
	svc.Labels = map[string]string{
		RabbitmqCommonLabelKey: RabbitmqCommonLabelValue,
	}
	svc.Spec = corev1.ServiceSpec{
		Selector: map[string]string{
			RabbitmqLabelKey: cluster.Name,
		},
		Type: corev1.ServiceTypeNodePort,
		Ports: []corev1.ServicePort{
			corev1.ServicePort{
				Name: "tcp",
				Port: 5672,
				TargetPort: intstr.IntOrString{
					IntVal: 5672,
				},
			},
			corev1.ServicePort{
				Name: "web",
				Port: 15672,
				TargetPort: intstr.IntOrString{
					IntVal: 15672,
				},
				NodePort: 31672,
			},
		},
	}
}

func MutateConfigMap(cluster *v1beta1.Svc, cm *corev1.ConfigMap) {
	cm.Labels = map[string]string{
		"addonmanager.kubernetes.io/mode": "Reconcile",
	}
	cm.Data = map[string]string{
		"enabled_plugins": "[rabbitmq_management,rabbitmq_peer_discovery_k8s].",
		"rabbitmq.conf":   "      loopback_users.guest = false\n\n      default_user = RABBITMQ_USER\n      default_pass = RABBITMQ_PASS\n      ## Clustering\n      cluster_formation.peer_discovery_backend = rabbit_peer_discovery_k8s\n      cluster_formation.k8s.host = kubernetes.default.svc.cluster.local\n      cluster_formation.k8s.address_type = hostname\n      #################################################\n      # public-service is rabbitmq-cluster's namespace#\n      #################################################\n      cluster_formation.k8s.hostname_suffix = .rabbitmq.default.svc.cluster.local\n      cluster_formation.node_cleanup.interval = 10\n      cluster_formation.node_cleanup.only_log_warning = true\n      cluster_partition_handling = autoheal\n      ## queue master locator\n      queue_master_locator=min-masters\n"}
}
